#include "DeadCodeEli.h"
#include "RDominateTree.h"

#include <set>
#include <map>
#include <vector>
bool DeadCodeEli::is_critical(Instruction *inst)
{ //call语句所对应的函数可能存在store指令，也为关键语句
    if (inst->isTerminator() || inst->is_store() || inst->is_call())
        return true;
    else
        return false;
}

void DeadCodeEli::mark(Function *func)
{
    std::vector<Instruction *> work_list;
    for (auto bb : func->get_basic_blocks()) //初始化不可删除的代码
    {
        for (auto inst : bb->get_instructions())
        {
            if (is_critical(inst))
            {
                inst_mark[inst] = true;
                work_list.push_back(inst);
            }
            else
            {
                inst_mark[inst] = false;
            }
        }
    }
    while (!work_list.empty())
    {
        auto inst = work_list.back();
        work_list.pop_back();
        for (auto operand : inst->get_operands()) //每一个操作数对应的指令都是critical的
        {
            auto op_var = dynamic_cast<Instruction *>(operand);
            if (op_var != nullptr && inst_mark[op_var] == false) //如果操作数为一个非mark的变量
            {
                inst_mark[op_var] = true;
                work_list.push_back(op_var);
            }
            auto jmp_locate = dynamic_cast<BasicBlock *>(operand);
            if (jmp_locate != nullptr) //如果操作数为一个基本快（在指令中含有跳转指令时会发生这种情况）
            {
                auto new_terminator = jmp_locate->get_terminator();
                if (inst_mark[new_terminator] == false)
                {
                    inst_mark[new_terminator] = true;
                    work_list.push_back(new_terminator);
                }
            }
        }
        for (auto reverse_dom_froniter : inst->get_parent()->get_rdom_frontier())
        { //逆向支配边界的的终止节点，如果是br指令的话，需要mark
            auto new_terminator = reverse_dom_froniter->get_terminator();
            if (new_terminator->is_br() && inst_mark[new_terminator] == false)
            {
                inst_mark[new_terminator] = true;
                work_list.push_back(new_terminator);
            }
        }
    }
}

BasicBlock *DeadCodeEli::nearest_marked_postdominator(Instruction *inst)
{ //采用广度优先算法，找出最近的被mark的后继支配节点
    std::list<BasicBlock *> bb_list;
    std::map<BasicBlock *, bool> is_bb_visited;
    auto inst_bb = inst->get_parent();
    auto inst_rdom = inst_bb->get_rdoms();
    for (auto bb : inst_rdom)
    {
        is_bb_visited[bb] = false;
    }
    bb_list.push_back(inst_bb);
    while (!bb_list.empty())
    {
        auto temp_bb = bb_list.front();
        bb_list.pop_front();
        if (is_bb_visited[temp_bb] == false)
        {
            is_bb_visited[temp_bb] = true;
            for (auto succ_bb : temp_bb->get_succ_basic_blocks())
            {
                if (succ_bb != inst_bb && inst_mark[succ_bb->get_terminator()] == true && inst_rdom.find(succ_bb) != inst_rdom.end())
                {
                    return succ_bb;
                }
                else if (is_bb_visited[succ_bb] == false)
                {
                    bb_list.push_back(succ_bb);
                }
            }
        }
    }
    return nullptr;
}

void DeadCodeEli::sweep(Function *func)
{
    std::vector<Instruction *> useless_inst;
    for (auto bb : func->get_basic_blocks()) //初始化不可删除的代码
    {
        for (auto inst : bb->get_instructions())
        {
            if (inst_mark[inst] == false)
            {
                if (inst->is_br()) //rewrite i with a jump to i’s nearest marked postdominator
                {
                    auto postdominator = nearest_marked_postdominator(inst);
                    inst->remove_use_of_ops();
                    inst->set_operand(0, postdominator);
                }
                else
                {
                    useless_inst.push_back(inst);
                }
            }
        }
        for (auto inst : useless_inst)
        {
            bb->delete_instr(inst);
        }
    }
}

void DeadCodeEli::execute()
{
    for (auto func : module->get_functions())
    {
        inst_mark.clear();
        RDominateTree reverse_dominate_tree(module);
        reverse_dominate_tree.execute();
        mark(func);
        sweep(func);
    }
    // Start from here!
}
